#undef UNICODE
#undef _UNICODE
#include "sym.h"


PVOID getModbase()
{
	NativeApi();
    NTSTATUS ntstatus;
    ULONG len =0;
    ntstatus = ZwQuerySystemInformation(SystemModuleInformation,NULL,0,&len);

    char* buff = NULL;
	/*
    if( !NT_SUCCESS( ntstatus ))
    {
        printf("[ZwQuerySystemInformation] error\n");
        printf("[ZwQuerySystemInformation] ntstatus = 0x%x\n", ntstatus);
        ;
    }*/
    buff = new char[len];
    memset(buff,0x00,len);
    ntstatus = ZwQuerySystemInformation(SystemModuleInformation,buff,len,&len);


    PSYSTEM_MODULE_INFORMATION psmi = (PSYSTEM_MODULE_INFORMATION)buff;

    return psmi->Module[0].Base;

}

ProcedureAddrRetrieve::ProcedureAddrRetrieve(PTCHAR PdbSearchPath,DWORD _Options )
{
    BOOL bRet = FALSE;
    ZeroMemory(this,sizeof(ProcedureAddrRetrieve));
    // Set options
    DWORD Options = SymGetOptions();

    // SYMOPT_DEBUG option asks DbgHelp to print additional troubleshooting
    // messages to debug output - use the debugger's Debug Output window
    // to view the messages

    Options =Options | SYMOPT_DEBUG | _Options;//|SYMOPT_ALLOW_ABSOLUTE_SYMBOLS

    SymSetOptions( Options );//0x201680
    // Initialize DbgHelp and load symbols for all modules of the current process

    bRet = SymInitialize (
               GetCurrentProcess(),  // Process handle of the current process
               PdbSearchPath,                 // user-defined search path -> use default
               FALSE                 // Do not load symbols for modules in the current process
           );

    SymSetSearchPath(GetCurrentProcess(),PdbSearchPath);
    if( !bRet )
    {
        _tprintf(_T("Error: SymInitialize() failed. Error code: %u \n"), ::GetLastError());
        return ;
    }
}
BOOL ProcedureAddrRetrieve::LoadSymbol(TCHAR *DllName,DWORD64 ModBase2)
{
    HMODULE ModBase;
    TCHAR DllFullPath[MAX_PATH];
    DWORD     FileSize  = 0;

    //Get DLL full path.
    ModBase = GetModuleHandle(DllName);
    if(ModBase == NULL)
    {
        ModBase = LoadLibrary(DllName);
        if(ModBase == NULL)
            return FALSE;
    }
    GetModuleFileName(ModBase,DllFullPath,MAX_PATH);

    GetFileSize(DllFullPath,FileSize);

    if(ModBase64)
    {
        // Unload symbols for the module
        if(!SymUnloadModule64( GetCurrentProcess(), ModBase64 ))
        {
            _tprintf( _T("Error: SymUnloadModule64() failed. Error code: %u \n"), ::GetLastError() );
            return FALSE;
        }
    }

    // Load symbols for the module
    _tprintf( _T("Loading symbols for: %s ... \n"), DllFullPath );


    ModBase64 = SymLoadModule64 (
                    GetCurrentProcess(), // Process handle of the current process
                    NULL,                // Handle to the module's image file (not needed)
                    DllFullPath,           // Path/name of the file
                    NULL,                // User-defined short name of the module (it can be NULL)
                    (DWORD64)ModBase,     //ModBase       // Base address of the module (cannot be NULL if .PDB file is used, otherwise it can be NULL)
                    FileSize             // Size of the file (cannot be NULL if .PDB file is used, otherwise it can be NULL)
                );

    /*
        ModBase64 = SymLoadModule64 (
                GetCurrentProcess(), // Process handle of the current process
                NULL,                // Handle to the module's image file (not needed)
                DllFullPath,           // Path/name of the file
                NULL,                // User-defined short name of the module (it can be NULL)
                NULL,            // Base address of the module (cannot be NULL if .PDB file is used, otherwise it can be NULL)
                NULL             // Size of the file (cannot be NULL if .PDB file is used, otherwise it can be NULL)
            );*/

    if( ModBase == 0 )
    {
        _tprintf(_T("Error: SymLoadModule64() failed. Error code: %u \n"), ::GetLastError());
        return FALSE;
    }
    return TRUE;
}
BOOL ProcedureAddrRetrieve::GetFileSize( const TCHAR* pFileName, DWORD& FileSize )
{
    // Check parameters
    if( pFileName == 0 )
    {
        return FALSE;
    }

    // Open the file
    HANDLE hFile = CreateFile( pFileName, GENERIC_READ, FILE_SHARE_READ,
                               NULL, OPEN_EXISTING, 0, NULL );

    if( hFile == INVALID_HANDLE_VALUE )
    {
        _tprintf( _T("CreateFile() failed. Error: %u \n"), ::GetLastError() );
        return FALSE;
    }

    // Obtain the size of the file
    FileSize = ::GetFileSize( hFile, NULL );
    if( FileSize == INVALID_FILE_SIZE )
    {
        _tprintf( _T("GetFileSize() failed. Error: %u \n"), ::GetLastError() );
        // and continue ...
    }

    // Close the file
    if( !CloseHandle( hFile ) )
    {
        _tprintf( _T("CloseHandle() failed. Error: %u \n"), ::GetLastError() );
        // and continue ...
    }
    // Complete

    return ( FileSize != INVALID_FILE_SIZE );

}
ProcedureAddrRetrieve::~ProcedureAddrRetrieve()
{
    if(ModBase64)
    {
        // Unload symbols for the module
        if(!SymUnloadModule64( GetCurrentProcess(), ModBase64 ))
        {
            _tprintf( _T("Error: SymUnloadModule64() failed. Error code: %u \n"), ::GetLastError() );
        }
    }
}
BOOL ProcedureAddrRetrieve::EnumSymbol(TCHAR *pSearchMask)
{
    BOOL bRet = TRUE;

    // Enumerate symbols and display information about them

    if( pSearchMask != NULL )
        _tprintf( _T("Search mask: %s \n"), pSearchMask );

    _tprintf( _T("Symbols: \n") );

    bRet = ::SymEnumSymbols(
               GetCurrentProcess(),   // Process handle of the current process
               ModBase64,               // Base address of the module
               pSearchMask,           // Mask (NULL -> all symbols)
               EnumSymbolsCallback, // The callback function
               this                   // A used-defined context can be passed here, if necessary
           );

    if( !bRet )
    {
        _tprintf( _T("Error: SymEnumSymbols() failed. Error code: %u \n"), ::GetLastError() );
    }
    return TRUE;
}
BOOL CALLBACK ProcedureAddrRetrieve::EnumSymbolsCallback( SYMBOL_INFO* pSymInfo, ULONG SymbolSize, PVOID UserContext )
{
    ProcedureAddrRetrieve *pPar = (ProcedureAddrRetrieve *)UserContext;
    if( pSymInfo != 0 )
    {
        if(*pPar->szProcedureName)
        {
            if(strcmp(pPar->szProcedureName,pSymInfo->Name)==0)
            {
                pPar->Address = pSymInfo->Address;
                return FALSE;
            }
        }
        else
        {
            _tprintf(pSymInfo->Name);
            _tprintf("\n");
        }
    }
    return TRUE;
}
DWORD64 ProcedureAddrRetrieve::RetrieveAddr(PTCHAR _szProcedureName)
{
    strncpy(szProcedureName,_szProcedureName,MAX_PATH);
    Address = 0;
    if(!EnumSymbol(""))
        return NULL;
    return Address;
}


ZWQuerySystemInformation ZwQuerySystemInformation;
bool NativeApi()
{
    HMODULE hNtDll;
    hNtDll = LoadLibrary( "ntdll.dll" );
    if(!hNtDll)
    {
        printf( "LoadLibrary( NTDLL.DLL ) Error:%d\n", GetLastError() );
        return false;
    }

 

	ZwQuerySystemInformation  = (ZWQuerySystemInformation)GetProcAddress(hNtDll, "ZwQuerySystemInformation");
    return true;
}